
import random
from cv2 import IMREAD_GRAYSCALE, rotate
import torch 
from PIL import Image
from torch.utils import data 
from torchvision import transforms
import os.path as osp
import numpy as np
from scipy.ndimage import zoom
from scipy.io import loadmat
import cv2
import os

'''
description: 
param {*} 
root: root dir of the data set. 
split: train or valid or test subdir
augment: True or False to indict whether to use data augmentation
return {*}
'''
class myDataset(data.Dataset):
    def __init__(
        self,
        root,
        split,
        augment = True,
    ):
        self.root = root
        self.split = split
        self.augment = augment
        self.ids = []
        
        self._set_files()
        
        
    def _set_files(self):
        self.img_dir = osp.join(self.root,'img')
        self.deform_dir = osp.join(self.root,'deformation_field')
        file_list = osp.join(
            self.root,'ids',self.split +'.txt'
        )

        with open(file_list) as f:
            contents = f.readlines()
        self.ids = [n[:-1] for n in contents]
    
    def rotate_img(self,image, angle, center=None, scale=1.0, borderValue = 0):
        # grab the dimensions of the image
        (h, w) = image.shape[:2]

        # if the center is None, initialize it as the center of
        # the image
        if center is None:
            center = (w // 2, h // 2)

        # perform the rotation
        M = cv2.getRotationMatrix2D(center, angle, scale)
        rotated = cv2.warpAffine(image, M, (w, h), borderValue = borderValue)

        # return the rotated image
        return rotated
    def _augmentation(self,Img,deformation,mask):
        """
        > The function takes in an image, a deformation field and a mask, and returns an augmented image,
        deformation field, and mask.
        augmentation including gray scale augmentation, yaw augmentation(rotation), multi-scale resize, 
        random translate and random noise.
        
        :param Img: the image to be augmented
        :param deformation: the deformation field to be augmented   
        :return: the augmented image and the deformation field
        """

        #### gray scale augmentation
        gray_scale = np.random.uniform(0.5,2)
        Img = Img**gray_scale
        
        # #multi-scale resize
        # size_scale = np.random.uniform(0.8,1.2)
        # Img = cv2.resize(Img,dsize = None,fx = size_scale,fy = size_scale,interpolation = cv2.INTER_CUBIC)
        # deformation = cv2.resize(deformation,dsize = None,fx = size_scale,fy = size_scale,interpolation = cv2.INTER_NEAREST)
        

        ### rotation augmentation  
        angle = np.random.uniform(-30,30)
        Img = self.rotate_img(Img,angle,borderValue = 1)
        deformation = self.rotate_img(deformation,angle)
        mask = self.rotate_img(mask,angle,borderValue = 0)
        angle = angle/180*np.pi
        x = deformation[:,:,0]*np.cos(angle) - deformation[:,:,1]*np.sin(angle)
        y = deformation[:,:,0]*np.sin(angle) + deformation[:,:,1]*np.cos(angle)
        deformation[:,:,0] = x
        deformation[:,:,1] = y
        

        ### random translate
        trans_h= np.random.randint(-20,20)
        trans_w= np.random.randint(-20,20)
        M = np.float32([[1,0,trans_h],[0,1,trans_w]])
        Img = cv2.warpAffine(Img,M,Img.shape[:2],borderValue=1)
        deformation = cv2.warpAffine(deformation,M,Img.shape[:2],borderValue=0)
        mask = cv2.warpAffine(mask,M,Img.shape[:2],borderValue=0)
        
        
        # random gaussian noise added on the image
        standard = np.random.uniform(0,0.03)
        noise_img= np.random.normal(0,standard,Img.shape)
        Img += noise_img
        Img = np.clip(Img,0,1)

        # random gaussian blur added on the image
        blur = np.random.uniform(0,1)
        if blur > 0.5:
            blur_size = np.random.randint(1,3)*2+1
            Img = cv2.GaussianBlur(Img,(blur_size,blur_size),0)



        return Img, deformation,mask

    def __getitem__(self, index):
        id = self.ids[index]
        # print(id)
        image_format = '.png'
        data_format = '.npy'
        img_path = osp.join(self.img_dir, id + image_format)
        img = cv2.imread(img_path)
        deform_path = osp.join(self.deform_dir,id+data_format)
        deformation = np.load(deform_path)
        Img = img.astype(np.float32)/255.
        mask = (Img<1).astype(np.uint8)
        if self.augment:
            # print('yes')
            Img, deformation,mask = self._augmentation(Img,deformation,mask)
        
    
        # 形态学处理
        kernel = np.ones((3,3),np.uint8)
        mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
        Img = np.transpose(Img,(2,0,1))
        mask = np.transpose(mask,(2,0,1))
        deformation = np.transpose(deformation,(2,0,1))
        return id,Img.astype(np.float32),deformation.astype(np.float32),mask.astype(np.float32)
    def __len__(self):
        return len(self.ids)
    
        
    
